import streamlit as st
import pickle
import faiss
import os
os.environ["KMP_DUPLICATE_LIB_OK"]="TRUE"
import torch
import torchvision.transforms as T
from PIL import Image

patch_h = 28
patch_w = 28
feat_dim = 384

# 设置页面标题
st.set_page_config(page_title="图片检索")

# 显示上传图片的输入框  
uploaded_file = st.file_uploader("上传图片", type=['jpg', 'png'])

transform = T.Compose([
    T.GaussianBlur(9, sigma=(0.1, 2.0)),
    T.Resize((patch_h * 14, patch_w * 14)),
    T.CenterCrop((patch_h * 14, patch_w * 14)),
    T.ToTensor(),
    T.Normalize(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)),
])

dinov2 = torch.hub.load('./model', 'dinov2_vits14', source='local').cuda()

# 加载保存的特征矩阵和URL映射表
with open('./faiss/feature_matrix.pkl', 'rb') as f:
    feature_matrix = pickle.load(f)

with open('./faiss/url_mapping.pkl', 'rb') as f:
    url_mapping = pickle.load(f)

# 创建Faiss索引
index = faiss.IndexFlatIP(feature_matrix.shape[1])
index.add(feature_matrix)


def get_images(image):
    global dinov2
    # image = Image.open(image_path)
    t_image = transform(image).unsqueeze(dim=0)
    querry_feature = dinov2(t_image.cuda())
    normalized_feature = querry_feature / querry_feature.norm(dim=-1, keepdim=True)
    normalized_feature = normalized_feature.detach().cpu().numpy()

    # 执行查询
    k = 6  # 要检索的最相似图片数量
    D, I = index.search(normalized_feature, k)

    # 获取最相似特征向量对应的图片URL
    similar_images = []
    print("最相似图片的URL：")
    for i in range(k):
        similar_image_index = I[0][i]
        similar_image_url = url_mapping[similar_image_index]
        print(similar_image_url)
        similar_images.append((Image.open(similar_image_url), similar_image_url))

    return similar_images


if uploaded_file is not None:
    # 加载上传的图片  
    image = Image.open(uploaded_file)
    st.image(image, caption=f"输入的图片")
    # 调用你的函数生成五张图片  
    images = get_images(image)

    # 定义图片大小
    img_size = (100, 100)  # 调整为你想要的大小
    # 九宫格布局
    col_width = st.columns(3)

    # 显示返回的图片  
    for i, img in enumerate(images):
        with col_width[i % 3]:
            st.image(img[0], caption=f"返回的图片{i + 1}     图片集路径：{img[1]}", use_column_width=True, width=200)  # 设置图片大小并按照九宫格展示